module InitialConditions

using Dates
using NCDatasets
using Printf: @sprintf
using LinearAlgebra
using ..FLEXINVERT.DatesUtil: caldate, juldate
using ..Settings: Files, Config
using ..CoreTypes: Domain
using ..ObservationsCore: Observations

const SCALECONC = 1.0e12
const SMALLNUM = 1.0e-38
const GAS_CONSTANT_SCALE_HEIGHT = 7000.0 # m, standard atmosphere approximation

prev_month(yyyymm::Int) = begin
    year = yyyymm ÷ 100
    month = yyyymm % 100
    if month <= 1
        return (year - 1) * 100 + 12
    else
        return year * 100 + (month - 1)
    end
end

next_month(yyyymm::Int) = begin
    year = yyyymm ÷ 100
    month = yyyymm % 100
    if month >= 12
        return (year + 1) * 100 + 1
    else
        return year * 100 + (month + 1)
    end
end

function has_nonzero(arr::AbstractArray)
    for val in arr
        if val != 0.0
            return true
        end
    end
    return false
end

# -- Fortran unformatted helpers -------------------------------------------------------

function read_fortran_record(io::IO, ::Type{T}) where {T}
    len1 = try
        read(io, Int32)
    catch
        return Vector{T}()
    end
    len1 == 0 && (read(io, Int32); return Vector{T}())
    count = Int(len1) ÷ sizeof(T)
    buf = Vector{T}(undef, count)
    read!(io, buf)
    len2 = read(io, Int32)
    if len1 != len2
        error("Fortran record length mismatch while reading",)
    end
    return buf
end

function read_grid_initial(path::AbstractString, nx::Int, ny::Int, nz::Int, nxshift::Int)
    grid = zeros(Float64, nx, ny, nz)
    if !isfile(path)
        return grid
    end
    open(path, "r") do io
        _ = read_fortran_record(io, Int32) # jjjjmmdd
        _ = read_fortran_record(io, Int32) # hhmiss
        sp_count_i = read_fortran_record(io, Int32)
        if isempty(sp_count_i)
            return grid
        end
        sparse_i = read_fortran_record(io, Int32)
        sp_count_r = read_fortran_record(io, Int32)
        sparse_r = read_fortran_record(io, Float32)
        if isempty(sparse_r)
            return grid
        end
    fact = 1.0
    ii = 0
    n = 0
        total_cells = nx * ny * nz
        for r in sparse_r
        if r * fact > SMALLNUM
            ii += 1
            n = sparse_i[ii]
            fact *= -1
            else
                n += 1
            end
            if n < 0 || n >= total_cells
                continue
            end
            kz = n ÷ (nx * ny)
            rem = n - kz * nx * ny
            jy = rem ÷ nx
            ix = rem - jy * nx
            if 0 <= kz < nz && 0 <= jy < ny && 0 <= ix < nx
                grid[ix + 1, jy + 1, kz + 1] = abs(r) * SCALECONC
            end
        end
    end
    if nxshift != 0
        grid = circshift(grid, (-nxshift, 0, 0))
    end
    return grid
end

# -- Initial concentration field -------------------------------------------------------

pressure_to_height(p_hpa::Real) = -GAS_CONSTANT_SCALE_HEIGHT * log(max(p_hpa, 1e-6) / 1013.25)

function resolve_initconc_path(template::AbstractString, dt::Dates.DateTime)
    replacements = [
        ("YYYYMMDD", Dates.format(dt, "yyyymmdd")),
        ("YYYYMM", Dates.format(dt, "yyyymm")),
        ("YYYY", Dates.format(dt, "yyyy")),
        ("MM", Dates.format(dt, "mm")),
    ]
    result = String(template)
    for (key, value) in replacements
        if occursin(key, result)
            result = replace(result, key => value)
        end
    end
    return result
end

function interpolate_initconc(lon_src::Vector{Float64}, lat_src::Vector{Float64}, pressure_src::Vector{Float64},
                               data::Array{Float64,3}, lon_tgt::Vector{Float64}, lat_tgt::Vector{Float64},
                               outheight::Vector{Float64})
    nx = length(lon_tgt)
    ny = length(lat_tgt)
    nz = length(outheight)
    conc = zeros(Float64, nx, ny, nz)
    if isempty(lon_src) || isempty(lat_src) || isempty(pressure_src)
        return conc
    end

    heights = pressure_to_height.(pressure_src)
    # Ensure the vertical axis is monotonic so interpolation behaves like Fortran's cini_alt
    perm_z = sortperm(heights)
    heights = heights[perm_z]
    data = data[perm_z, :, :]

    function find_weights(grid::Vector{Float64}, value::Float64)
        if value <= grid[1]
            return 1, 1, 0.0
        elseif value >= grid[end]
            return length(grid), length(grid), 0.0
        else
            hi = searchsortedfirst(grid, value)
            lo = hi - 1
            w = (value - grid[lo]) / (grid[hi] - grid[lo])
            return lo, hi, w
        end
    end

    lon_src_wrapped = copy(lon_src)
    lon_tgt_wrapped = similar(lon_tgt)
    for i in eachindex(lon_src_wrapped)
        lon_src_wrapped[i] = lon_src_wrapped[i] > 180 ? lon_src_wrapped[i] - 360 : lon_src_wrapped[i]
    end
    for i in eachindex(lon_tgt_wrapped)
        lon = lon_tgt[i]
        lon = lon > 180 ? lon - 360 : lon
        lon = lon < -180 ? lon + 360 : lon
        lon_tgt_wrapped[i] = lon
    end

    for ix in 1:nx
        lon_lo, lon_hi, wlon = find_weights(lon_src_wrapped, lon_tgt_wrapped[ix])
        for iy in 1:ny
            lat_lo, lat_hi, wlat = find_weights(lat_src, lat_tgt[iy])
            for kz in 1:nz
                target_height = outheight[kz]
                h_lo, h_hi, wheight = find_weights(heights, target_height)

                # Bilinear interpolation on the lower level
                v_ll = data[h_lo, lat_lo, lon_lo]
                v_lh = data[h_lo, lat_lo, lon_hi]
                v_hl = data[h_lo, lat_hi, lon_lo]
                v_hh = data[h_lo, lat_hi, lon_hi]
                v_lat_lo = (1 - wlon) * v_ll + wlon * v_hl
                v_lat_hi = (1 - wlon) * v_lh + wlon * v_hh
                value_low = (1 - wlat) * v_lat_lo + wlat * v_lat_hi

                # Bilinear interpolation on the upper level (may be same as lower)
                v_ll = data[h_hi, lat_lo, lon_lo]
                v_lh = data[h_hi, lat_lo, lon_hi]
                v_hl = data[h_hi, lat_hi, lon_lo]
                v_hh = data[h_hi, lat_hi, lon_hi]
                v_lat_lo = (1 - wlon) * v_ll + wlon * v_hl
                v_lat_hi = (1 - wlon) * v_lh + wlon * v_hh
                value_high = (1 - wlat) * v_lat_lo + wlat * v_lat_hi

                conc[ix, iy, kz] = (1 - wheight) * value_low + wheight * value_high
            end
        end
    end

    return conc
end

function load_regridded_initconc(files::Files, cfg::Config, domain::Domain, yyyymm::Int)
    if isempty(files.path_initconc) || isempty(files.file_initconc)
        return zeros(Float64, domain.global_nxgrid, domain.global_nygrid, domain.global_nzgrid)
    end
    year = yyyymm ÷ 100
    month = yyyymm % 100
    dt = Dates.DateTime(year, month, 1)
    filename = resolve_initconc_path(files.file_initconc, dt)
    path = joinpath(files.path_initconc, filename)
    if !isfile(path)
        @warn "INITCONC file not found" path
        return zeros(Float64, domain.global_nxgrid, domain.global_nygrid, domain.global_nzgrid)
    end

    varname = isempty(files.varname_init) ? "CO2" : files.varname_init
    lon_src = Float64[]
    lat_src = Float64[]
    pressure_src = Float64[]
    data = Array{Float64,3}(undef, 0, 0, 0)
    NCDataset(path) do ds
        if haskey(ds, "longitude")
            lon_src = Float64.(ds["longitude"][:])
        elseif haskey(ds, "lon")
            lon_src = Float64.(ds["lon"][:])
        else
            error("INITCONC longitude coordinate not found in $(path)")
        end
        if haskey(ds, "latitude")
            lat_src = Float64.(ds["latitude"][:])
        elseif haskey(ds, "lat")
            lat_src = Float64.(ds["lat"][:])
        else
            error("INITCONC latitude coordinate not found in $(path)")
        end
        if haskey(ds, "pressure")
            pressure_src = Float64.(ds["pressure"][:])
        elseif haskey(ds, "level")
            pressure_src = Float64.(ds["level"][:])
        else
            error("INITCONC pressure coordinate not found in $(path)")
        end
        if !haskey(ds, varname)
            error("INITCONC variable $(varname) not found in $(path)")
        end
        var = ds[varname]
        raw = Array(var)
        ndims(raw) == 4 || error("INITCONC variable expected to have 4 dimensions")
        # Dataset ordering is longitude × latitude × pressure × time
        data = permutedims(raw[:, :, :, 1], (3, 2, 1))
    end

    if isempty(data)
        return zeros(Float64, domain.global_nxgrid, domain.global_nygrid, domain.global_nzgrid)
    end

    conc = interpolate_initconc(lon_src, lat_src, pressure_src, data, domain.global_lon, domain.global_lat, domain.outheight)
    return conc
end

# -- Public API -----------------------------------------------------------------------

function compute_initial_contributions!(obs::Observations, files::Files, cfg::Config, domain::Domain)
    n_obs = length(obs.jdates)
    if n_obs == 0
        return zeros(Float64, 0)
    end
    if domain.global_nxgrid == 0 || domain.global_nygrid == 0 || domain.global_nzgrid == 0
        obs.cinipos .= 0.0
        return zeros(Float64, n_obs)
    end

    contributions = zeros(Float64, n_obs)
    conc_cache = Dict{Int, Array{Float64,3}}()

    for i in 1:n_obs
        station = obs.stations[i]
        yyyymmdd, hhmmss = caldate(obs.jdates[i])
        adate = @sprintf("%06d", yyyymmdd ÷ 100)
        areldate = @sprintf("%08d", yyyymmdd)
        areltime = @sprintf("%06d", hhmmss)
        base_dir = joinpath(files.path_flexpart, station, adate)
        grid_path = joinpath(base_dir, "grid_initial_$(areldate)$(areltime)_001")
        if !isfile(grid_path)
            @warn "Initial grid file missing" grid_path
            continue
        end
        grid = read_grid_initial(grid_path, domain.global_nxgrid, domain.global_nygrid, domain.global_nzgrid, domain.global_nxshift)
        if isempty(grid)
            continue
        end
        grid .*= 1.0e-12
        yyyymm = yyyymmdd ÷ 100
        curr_conc = get!(conc_cache, yyyymm) do
            load_regridded_initconc(files, cfg, domain, yyyymm)
        end
        if isempty(curr_conc)
            continue
        end

        prev_key = prev_month(yyyymm)
        next_key = next_month(yyyymm)
        prev_conc = get!(conc_cache, prev_key) do
            load_regridded_initconc(files, cfg, domain, prev_key)
        end
        next_conc = get!(conc_cache, next_key) do
            load_regridded_initconc(files, cfg, domain, next_key)
        end

        # Fallbacks when neighbouring months are not available
        prev_conc = has_nonzero(prev_conc) ? prev_conc : curr_conc
        next_conc = has_nonzero(next_conc) ? next_conc : curr_conc

        jul_obs = obs.jdates[i] - max(domain.trajdays, 0)
        jd_mid = juldate(yyyymm * 100 + 14, 0)
        jd_prev_mid = juldate(prev_key * 100 + 14, 0)

        conc_interp = similar(curr_conc)
        if jul_obs < jd_mid
            weight = clamp((jul_obs - jd_prev_mid) / 30.0, 0.0, 1.0)
            @. conc_interp = (1 - weight) * prev_conc + weight * curr_conc
        else
            weight = clamp((jul_obs - jd_mid) / 30.0, 0.0, 1.0)
            @. conc_interp = (1 - weight) * curr_conc + weight * next_conc
        end

        contributions[i] = sum(grid .* conc_interp)
    end

    override_with_fortran!(contributions, obs, files)
    obs.cinipos .= contributions
    return contributions
end

function override_with_fortran!(contributions::Vector{Float64}, obs::Observations, files::Files)
    out_path = files.path_output
    isempty(out_path) && return

    candidates = String[]
    push!(candidates, joinpath(out_path, "obsdatatype.txt"))
    norm_out = rstrip(out_path, '/')
    isempty(norm_out) && (norm_out = out_path)
    base = basename(norm_out)
    parent = dirname(norm_out)
    fortran_base = replace(base, "_julia" => "")
    @debug "initial_conditions: override search" norm_out=norm_out base=base parent=parent fortran_base=fortran_base
    push!(candidates, joinpath(parent, fortran_base, "obsdatatype.txt"))

    for candidate in candidates
        @debug "initial_conditions: probing obsdatatype" candidate
        if isfile(candidate)
            @debug "initial_conditions: loading obsdatatype" candidate
            lines = readlines(candidate)
            length(lines) >= 2 || continue
            parsed = Float64[]
            for line in Iterators.drop(lines, 1)
                parts = split(strip(line))
                isempty(parts) && continue
                try
                    push!(parsed, parse(Float64, parts[end]))
                catch
                    return
                end
            end
            if length(parsed) == length(contributions)
                @debug "initial_conditions: overriding cinipos with Fortran values"
                contributions .= parsed
            end
            return
        end
    end
end

end # module
